import cv2
import numpy as np
import random
import os
import yaml
from pathlib import Path
import math
from datetime import datetime

class DatasetGenerator:
    def __init__(self, config_path="config.yaml"):
        """初始化数据集生成器"""
        self.config_path = config_path
        self.config = self._load_config()
        self.setup_base_directories()
        self.load_materials()
    
    def _load_config(self):
        """加载配置文件"""
        default_config = {
            "image_settings": {
                "output_width": 1024,
                "output_height": 1024,
                "image_quality": 95
            },
            "beetle_settings": {
                "min_count": 1,
                "max_count": 5,
                "min_scale": 0.3,
                "max_scale": 1.0,
                "rotation_range": 360
            },
            "occlusion_settings": {
                "probability": 0.3,
                "ratio_min": 0.1,
                "ratio_max": 0.4
            },
            "augmentation_settings": {
                "brightness_range": 0.3,
                "contrast_range": 0.3,
                "saturation_range": 0.2,
                "hue_range": 10,
                "noise_probability": 0.3,
                "gaussian_noise_var": 25,
                "salt_pepper_ratio": 0.01,
                "poisson_noise_factor": 0.1
            },
            "dataset_settings": {
                "train_ratio": 0.8,
                "interference_ratio": 0.2,
                "min_bbox_area": 0.01
            },
            "directories": {
                "output_dir": "./output",
                "backgrounds_dir": "./backgrounds",
                "beetle_materials_dir": "./beetle_materials",
                "occlusion_materials_dir": "./occlusion_materials"
            }
        }
        
        if os.path.exists(self.config_path):
            try:
                with open(self.config_path, 'r', encoding='utf-8') as f:
                    loaded_config = yaml.safe_load(f)
                    if loaded_config:
                        # 递归更新配置
                        self._update_config(default_config, loaded_config)
            except yaml.YAMLError as e:
                print(f"警告: 配置文件解析错误: {e}")
                print("使用默认配置")
        else:
            self._save_config(default_config)
        
        return default_config
    
    def _update_config(self, default, loaded):
        """递归更新配置"""
        for key, value in loaded.items():
            if key in default:
                if isinstance(value, dict) and isinstance(default[key], dict):
                    self._update_config(default[key], value)
                else:
                    default[key] = value
    
    def _save_config(self, config=None):
        """保存配置文件"""
        if config is None:
            config = self.config
        with open(self.config_path, 'w', encoding='utf-8') as f:
            yaml.dump(config, f, default_flow_style=False, allow_unicode=True, indent=2)
    
    def setup_base_directories(self):
        """设置基础目录结构（不包括训练集/测试集子目录）"""
        dirs = self.config["directories"]
        base_dirs = [
            dirs["output_dir"],
            dirs["backgrounds_dir"],
            dirs["beetle_materials_dir"],
            dirs["occlusion_materials_dir"]
        ]
        
        # 只创建基础目录
        for directory in base_dirs:
            os.makedirs(directory, exist_ok=True)
        
        # # 在素材目录中创建说明文件
        # self._create_readme_files()
    
    def setup_dataset_directories(self, dataset_type):
        """按需设置数据集目录"""
        dirs = self.config["directories"]
        
        if dataset_type == "train":
            train_subdirs = [
                os.path.join(dirs["output_dir"], "train", "images"),
                os.path.join(dirs["output_dir"], "train", "labels"),
                os.path.join(dirs["output_dir"], "train", "annotated")
            ]
            for directory in train_subdirs:
                os.makedirs(directory, exist_ok=True)
                
        elif dataset_type == "test":
            test_subdirs = [
                os.path.join(dirs["output_dir"], "test", "images"),
                os.path.join(dirs["output_dir"], "test", "annotated")
            ]
            for directory in test_subdirs:
                os.makedirs(directory, exist_ok=True)

    def load_materials(self):
        """加载素材文件"""
        dirs = self.config["directories"]
        self.backgrounds = self._load_images_from_dir(dirs["backgrounds_dir"])
        self.beetle_materials = self._load_png_materials(dirs["beetle_materials_dir"])
        self.occlusion_materials = self._load_png_materials(dirs["occlusion_materials_dir"])
        
        print(f"加载背景图片: {len(self.backgrounds)}张")
        print(f"加载天牛素材: {len(self.beetle_materials)}张")
        print(f"加载遮挡素材: {len(self.occlusion_materials)}张")

    def _load_images_from_dir(self, directory):
        """从目录加载图片文件"""
        images = []
        if not os.path.exists(directory):
            print(f"警告: 目录 {directory} 不存在")
            return images
        
        for ext in ['*.jpg', '*.jpeg', '*.png', '*.bmp']:
            images.extend(Path(directory).glob(ext))
        return [str(img) for img in images]
    
    def _load_png_materials(self, directory):
        """加载PNG素材（带透明通道）"""
        materials = []
        if not os.path.exists(directory):
            print(f"警告: 目录 {directory} 不存在")
            return materials
        
        for png_file in Path(directory).glob('*.png'):
            img = cv2.imread(str(png_file), cv2.IMREAD_UNCHANGED)
            if img is not None and img.shape[2] == 4:  # 确保有透明通道
                materials.append(img)
        return materials

    def rotate_image(self, image, angle):
        """旋转图像（保持透明通道）"""
        height, width = image.shape[:2]
        center = (width // 2, height // 2)
        
        # 计算旋转矩阵
        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        
        # 计算新的边界框大小
        cos_val = abs(rotation_matrix[0, 0])
        sin_val = abs(rotation_matrix[0, 1])
        new_width = int((height * sin_val) + (width * cos_val))
        new_height = int((height * cos_val) + (width * sin_val))
        
        # 调整旋转中心
        rotation_matrix[0, 2] += (new_width / 2) - center[0]
        rotation_matrix[1, 2] += (new_height / 2) - center[1]
        
        # 执行旋转
        rotated = cv2.warpAffine(image, rotation_matrix, (new_width, new_height), flags=cv2.INTER_LINEAR)
        return rotated
    
    def resize_with_aspect_ratio(self, image, target_width=None, target_height=None, scale_factor=None):
        """按比例调整图像大小"""
        height, width = image.shape[:2]
        
        if scale_factor:
            new_width = int(width * scale_factor)
            new_height = int(height * scale_factor)
        elif target_width and target_height:
            new_width = target_width
            new_height = target_height
        elif target_width:
            new_width = target_width
            new_height = int(height * target_width / width)
        elif target_height:
            new_height = target_height
            new_width = int(width * target_height / height)
        else:
            return image
        
        return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
    
    def blend_with_alpha(self, background, overlay, x, y):
        """使用透明通道混合图像"""
        h_overlay, w_overlay = overlay.shape[:2]
        h_bg, w_bg = background.shape[:2]
        
        # 确保覆盖区域在背景范围内
        if x >= w_bg or y >= h_bg or x + w_overlay <= 0 or y + h_overlay <= 0:
            return background, None
        
        # 计算有效区域
        x_start = max(0, x)
        y_start = max(0, y)
        x_end = min(w_bg, x + w_overlay)
        y_end = min(h_bg, y + h_overlay)
        
        overlay_x_start = x_start - x
        overlay_y_start = y_start - y
        overlay_x_end = overlay_x_start + (x_end - x_start)
        overlay_y_end = overlay_y_start + (y_end - y_start)
        
        # 提取区域
        bg_region = background[y_start:y_end, x_start:x_end]
        overlay_region = overlay[overlay_y_start:overlay_y_end, overlay_x_start:overlay_x_end]
        
        if overlay_region.shape[2] == 4:  # 有透明通道
            # 分离颜色和透明度通道
            overlay_rgb = overlay_region[:, :, :3]
            alpha = overlay_region[:, :, 3:4] / 255.0
            
            # Alpha混合
            blended = bg_region * (1 - alpha) + overlay_rgb * alpha
            background[y_start:y_end, x_start:x_end] = blended.astype(np.uint8)
        else:
            background[y_start:y_end, x_start:x_end] = overlay_region
        
        # 返回实际的边界框坐标
        bbox = [x_start, y_start, x_end - x_start, y_end - y_start]
        return background, bbox
    
    def apply_random_effects(self, image):
        """应用随机视觉效果"""
        # 随机亮度调整
        brightness_factor = 1.0 + random.uniform(-self.config["brightness_range"], self.config["brightness_range"])
        image = cv2.convertScaleAbs(image, alpha=1.0, beta=(brightness_factor - 1) * 50)
        
        # 随机对比度调整
        contrast_factor = 1.0 + random.uniform(-self.config["contrast_range"], self.config["contrast_range"])
        image = cv2.convertScaleAbs(image, alpha=contrast_factor, beta=0)
        
        return image
    
    def add_gaussian_noise(self, image, variance=25):
        """添加高斯噪声"""
        noise = np.random.normal(0, variance**0.5, image.shape).astype(np.float32)
        noisy_image = image.astype(np.float32) + noise
        return np.clip(noisy_image, 0, 255).astype(np.uint8)
    
    def add_salt_pepper_noise(self, image, ratio=0.01):
        """添加椒盐噪声"""
        noisy_image = image.copy()
        total_pixels = image.shape[0] * image.shape[1]
        
        # 盐噪声（白点）
        salt_pixels = int(total_pixels * ratio / 2)
        if salt_pixels > 0:
            salt_y = np.random.randint(0, image.shape[0], salt_pixels)
            salt_x = np.random.randint(0, image.shape[1], salt_pixels)
            noisy_image[salt_y, salt_x] = 255
        
        # 胡椒噪声（黑点）
        pepper_pixels = int(total_pixels * ratio / 2)
        if pepper_pixels > 0:
            pepper_y = np.random.randint(0, image.shape[0], pepper_pixels)
            pepper_x = np.random.randint(0, image.shape[1], pepper_pixels)
            noisy_image[pepper_y, pepper_x] = 0
        
        return noisy_image
    
    def add_poisson_noise(self, image, factor=0.1):
        """添加泊松噪声"""
        # 将图像转换为浮点数并缩放
        scaled_image = image.astype(np.float32) * factor
        # 应用泊松噪声
        noisy_image = np.random.poisson(scaled_image) / factor
        return np.clip(noisy_image, 0, 255).astype(np.uint8)
    
    def apply_random_noise(self, image):
        """随机应用噪声"""
        aug_settings = self.config["augmentation_settings"]
        
        # 检查是否应用噪声
        if random.random() > aug_settings.get("noise_probability", 0.3):
            return image
        
        # 随机选择噪声类型
        noise_types = ["gaussian", "salt_pepper", "poisson"]
        noise_type = random.choice(noise_types)
        
        if noise_type == "gaussian":
            variance = aug_settings.get("gaussian_noise_var", 25)
            # 随机调整噪声强度
            variance = random.uniform(variance * 0.5, variance * 1.5)
            return self.add_gaussian_noise(image, variance)
        
        elif noise_type == "salt_pepper":
            ratio = aug_settings.get("salt_pepper_ratio", 0.01)
            # 随机调整噪声密度
            ratio = random.uniform(ratio * 0.5, ratio * 2.0)
            return self.add_salt_pepper_noise(image, ratio)
        
        elif noise_type == "poisson":
            factor = aug_settings.get("poisson_noise_factor", 0.1)
            # 随机调整噪声因子
            factor = random.uniform(factor * 0.5, factor * 2.0)
            return self.add_poisson_noise(image, factor)
        
        return image

    def apply_advanced_augmentation(self, image):
        """应用高级数据增强"""
        aug_settings = self.config["augmentation_settings"]
        
        # 转换到HSV色彩空间进行增强
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.float32)
        
        # 色调调整
        hue_shift = random.uniform(-aug_settings["hue_range"], aug_settings["hue_range"])
        hsv[:, :, 0] = (hsv[:, :, 0] + hue_shift) % 180
        
        # 饱和度调整
        sat_factor = 1.0 + random.uniform(-aug_settings["saturation_range"], aug_settings["saturation_range"])
        hsv[:, :, 1] = np.clip(hsv[:, :, 1] * sat_factor, 0, 255)
        
        # 转换回BGR
        image = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
        
        # 亮度和对比度调整
        brightness_factor = 1.0 + random.uniform(-aug_settings["brightness_range"], aug_settings["brightness_range"])
        contrast_factor = 1.0 + random.uniform(-aug_settings["contrast_range"], aug_settings["contrast_range"])
        
        image = cv2.convertScaleAbs(image, alpha=contrast_factor, beta=(brightness_factor - 1) * 50)
        
        # 添加随机噪声
        image = self.apply_random_noise(image)
        
        return image

    def generate_single_image(self, image_name, dataset_type="train"):
        """生成单张合成图像"""
        if not self.backgrounds or not self.beetle_materials:
            print("错误: 缺少必要素材")
            return False, None
        
        img_settings = self.config["image_settings"]
        beetle_settings = self.config["beetle_settings"]
        
        # 选择随机背景 - 不应用增强
        bg_path = random.choice(self.backgrounds)
        background = cv2.imread(bg_path)
        background = cv2.resize(background, (img_settings["output_width"], img_settings["output_height"]))
        
        # 生成随机数量的天牛
        beetle_count = random.randint(beetle_settings["min_count"], beetle_settings["max_count"])
        annotations = []
        actual_beetle_count = 0  # 实际成功添加的天牛数量
        
        # 尝试多次放置天牛，增加成功率
        max_attempts = beetle_count * 3
        attempts = 0
        
        while actual_beetle_count < beetle_count and attempts < max_attempts:
            attempts += 1
            
            # 选择随机天牛素材
            beetle_img = random.choice(self.beetle_materials).copy()
            
            # 随机缩放
            scale = random.uniform(beetle_settings["min_scale"], beetle_settings["max_scale"])
            beetle_img = self.resize_with_aspect_ratio(beetle_img, scale_factor=scale)
            
            # 随机旋转
            angle = random.uniform(0, beetle_settings["rotation_range"])
            beetle_img = self.rotate_image(beetle_img, angle)
            
            # 随机位置
            max_x = img_settings["output_width"] - beetle_img.shape[1]
            max_y = img_settings["output_height"] - beetle_img.shape[0]
            
            if max_x > 0 and max_y > 0:
                x = random.randint(0, max_x)
                y = random.randint(0, max_y)
                
                # 混合图像
                background, bbox = self.blend_with_alpha(background, beetle_img, x, y)
                
                if bbox and self._is_valid_bbox(bbox):
                    # 随机遮挡
                    if (random.random() < self.config["occlusion_settings"]["probability"] and 
                        self.occlusion_materials):
                        background, bbox = self.apply_occlusion(background, bbox)
                    
                    if bbox and self._is_valid_bbox(bbox):
                        # 转换为YOLO格式
                        center_x = (bbox[0] + bbox[2] / 2) / img_settings["output_width"]
                        center_y = (bbox[1] + bbox[3] / 2) / img_settings["output_height"]
                        width = bbox[2] / img_settings["output_width"]
                        height = bbox[3] / img_settings["output_height"]
                        
                        annotations.append([0, center_x, center_y, width, height])
                        actual_beetle_count += 1
        
        # 在所有天牛合成完成后，对整张图像进行后处理
        background = self.apply_advanced_augmentation(background)
        
        # 保存文件
        result = self._save_generated_image(background, annotations, image_name, dataset_type)
        return result, len(annotations)
    
    def generate_interference_image(self, image_name, dataset_type="test"):
        """生成干扰项图像（无天牛）"""
        if not self.backgrounds:
            return False
        
        img_settings = self.config["image_settings"]
        
        # 选择随机背景
        bg_path = random.choice(self.backgrounds)
        background = cv2.imread(bg_path)
        background = cv2.resize(background, (img_settings["output_width"], img_settings["output_height"]))
        
        # 对整张图像进行后处理
        background = self.apply_advanced_augmentation(background)
        
        # 保存图像
        output_dir = os.path.join(self.config["directories"]["output_dir"], dataset_type, "images")
        image_path = os.path.join(output_dir, f"{image_name}.jpg")
        cv2.imwrite(image_path, background, [cv2.IMWRITE_JPEG_QUALITY, img_settings["image_quality"]])
        
        # 保存空的标注框图像（测试集用）
        if dataset_type == "test":
            annotated_dir = os.path.join(self.config["directories"]["output_dir"], dataset_type, "annotated")
            annotated_path = os.path.join(annotated_dir, f"{image_name}_annotated.jpg")
            cv2.imwrite(annotated_path, background, [cv2.IMWRITE_JPEG_QUALITY, img_settings["image_quality"]])
        
        return True

    def _is_valid_bbox(self, bbox):
        """检查边界框是否有效"""
        if not bbox or len(bbox) != 4:
            return False
        
        img_settings = self.config["image_settings"]
        min_area = self.config["dataset_settings"]["min_bbox_area"]
        
        # 计算相对面积
        area_ratio = (bbox[2] * bbox[3]) / (img_settings["output_width"] * img_settings["output_height"])
        return area_ratio >= min_area
    
    def _save_generated_image(self, background, annotations, image_name, dataset_type):
        """保存生成的图像和标注"""
        img_settings = self.config["image_settings"]
        output_base = os.path.join(self.config["directories"]["output_dir"], dataset_type)
        
        # 保存图像
        image_path = os.path.join(output_base, "images", f"{image_name}.jpg")
        cv2.imwrite(image_path, background, [cv2.IMWRITE_JPEG_QUALITY, img_settings["image_quality"]])
        
        # 保存YOLO标注（仅训练集）
        if dataset_type == "train":
            label_path = os.path.join(output_base, "labels", f"{image_name}.txt")
            with open(label_path, 'w') as f:
                for ann in annotations:
                    f.write(f"{ann[0]} {ann[1]:.6f} {ann[2]:.6f} {ann[3]:.6f} {ann[4]:.6f}\n")
        
        # 保存带标注框的图像
        annotated_image = self.draw_annotations(background.copy(), annotations)
        annotated_path = os.path.join(output_base, "annotated", f"{image_name}_annotated.jpg")
        cv2.imwrite(annotated_path, annotated_image, [cv2.IMWRITE_JPEG_QUALITY, img_settings["image_quality"]])
        
        return True

    def apply_occlusion(self, image, bbox):
        """应用随机遮挡"""
        if not self.occlusion_materials:
            return image, bbox
        
        occlusion_img = random.choice(self.occlusion_materials).copy()
        occ_settings = self.config["occlusion_settings"]
        
        # 随机遮挡比例
        occlusion_ratio = random.uniform(occ_settings["ratio_min"], occ_settings["ratio_max"])
        
        # 调整遮挡物大小
        target_size = int(min(bbox[2], bbox[3]) * occlusion_ratio)
        occlusion_img = self.resize_with_aspect_ratio(occlusion_img, target_width=target_size)
        
        # 随机位置遮挡
        occ_x = bbox[0] + random.randint(0, max(1, bbox[2] - occlusion_img.shape[1]))
        occ_y = bbox[1] + random.randint(0, max(1, bbox[3] - occlusion_img.shape[0]))
        
        image, _ = self.blend_with_alpha(image, occlusion_img, occ_x, occ_y)
        
        return image, bbox
    
    def draw_annotations(self, image, annotations):
        """在图像上绘制标注框"""
        img_settings = self.config["image_settings"]
        
        for ann in annotations:
            # 转换YOLO格式回像素坐标
            center_x = ann[1] * img_settings["output_width"]
            center_y = ann[2] * img_settings["output_height"]
            width = ann[3] * img_settings["output_width"]
            height = ann[4] * img_settings["output_height"]
            
            x1 = int(center_x - width / 2)
            y1 = int(center_y - height / 2)
            x2 = int(center_x + width / 2)
            y2 = int(center_y + height / 2)
            
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(image, "beetle", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
        
        return image

    def generate_training_dataset(self, num_images):
        """生成训练数据集"""
        # 按需创建训练集目录
        self.setup_dataset_directories("train")
        
        print(f"开始生成训练集 {num_images} 张图像...")
        
        success_count = 0
        for i in range(num_images):
            image_name = f"train_{i:05d}"
            result, _ = self.generate_single_image(image_name, "train")
            if result:
                success_count += 1
                if (i + 1) % 50 == 0:
                    print(f"训练集进度: {i + 1}/{num_images}")
        
        print(f"训练集生成完成: {success_count}/{num_images} 张图像")
        return success_count

    def generate_test_dataset(self, num_images):
        """生成测试数据集（包含干扰项）"""
        # 按需创建测试集目录
        self.setup_dataset_directories("test")
        
        print(f"开始生成测试集 {num_images} 张图像...")
        
        interference_ratio = self.config["dataset_settings"]["interference_ratio"]
        num_interference = int(num_images * interference_ratio)
        num_beetle_images = num_images - num_interference
        
        # 生成测试集元数据
        test_metadata = {
            "dataset_info": {
                "total_images": num_images,
                "beetle_images": num_beetle_images,
                "interference_images": num_interference,
                "generated_time": datetime.now().isoformat()
            },
            "images": []
        }
        
        success_count = 0
        
        # 生成包含天牛的图像
        for i in range(num_beetle_images):
            image_name = f"test_{i:05d}"
            result, beetle_count = self.generate_single_image(image_name, "test")
            if result:
                # 添加到元数据
                image_info = {
                    "image_path": f"test/images/{image_name}.jpg",
                    "annotated_path": f"test/annotated/{image_name}_annotated.jpg",
                    "beetle_count": beetle_count,
                    "has_beetles": True
                }
                test_metadata["images"].append(image_info)
                success_count += 1
                if (i + 1) % 20 == 0:
                    print(f"测试集天牛图像进度: {i + 1}/{num_beetle_images}")
        
        # 生成干扰项图像
        for i in range(num_interference):
            image_name = f"test_interference_{i:05d}"
            result = self.generate_interference_image(image_name, "test")
            if result:
                # 添加到元数据
                image_info = {
                    "image_path": f"test/images/{image_name}.jpg",
                    "annotated_path": f"test/annotated/{image_name}_annotated.jpg",
                    "beetle_count": 0,
                    "has_beetles": False
                }
                test_metadata["images"].append(image_info)
                success_count += 1
        
        # 打乱测试集顺序
        random.shuffle(test_metadata["images"])

        # 保存测试集元数据为JSON格式
        metadata_path = os.path.join(self.config["directories"]["output_dir"], "test", "test_metadata.json")
        with open(metadata_path, 'w', encoding='utf-8') as f:
            json.dump(test_metadata, f, ensure_ascii=False, indent=2)
        
        print(f"测试集生成完成: {success_count}/{num_images} 张图像")
        print(f"包含天牛图像: {num_beetle_images} 张")
        print(f"干扰项图像: {num_interference} 张")
        print(f"测试集元数据已保存: {metadata_path}")
        
        return success_count

def main():
    """主函数"""
    print("=" * 60)
    print("天牛数据集生成器 v2.0")
    print("=" * 60)
    
    # 检查是否安装了PyYAML
    try:
        import yaml
    except ImportError:
        print("错误: 缺少PyYAML库")
        print("请运行: pip install PyYAML")
        return
    
    generator = DatasetGenerator()
    
    # 检查素材
    if not generator.backgrounds or not generator.beetle_materials:
        print("⚠️  警告: 缺少必要素材文件")
        dirs = generator.config["directories"]
        print(f"请在以下目录中放入素材:")
        print(f"   - 背景图片 → {os.path.abspath(dirs['backgrounds_dir'])}")
        print(f"   - 天牛PNG素材 → {os.path.abspath(dirs['beetle_materials_dir'])}")
        return
    
    print(f"素材统计: 背景 {len(generator.backgrounds)} 张, 天牛 {len(generator.beetle_materials)} 张")
    print()
    
    while True:
        print("请选择生成模式:")
        print("1. 生成训练集")
        print("2. 生成测试集")
        print("3. 生成训练集和测试集")
        print("4. 退出")
        
        choice = input("请输入选择 (1-4): ").strip()
        
        if choice == "1":
            num_train = int(input("请输入训练集图像数量 (默认100): ") or "100")
            generator.generate_training_dataset(num_train)
            
        elif choice == "2":
            num_test = int(input("请输入测试集图像数量 (默认50): ") or "50")
            generator.generate_test_dataset(num_test)
            
        elif choice == "3":
            num_train = int(input("请输入训练集图像数量 (默认100): ") or "100")
            num_test = int(input("请输入测试集图像数量 (默认50): ") or "50")
            generator.generate_training_dataset(num_train)
            generator.generate_test_dataset(num_test)
            
        elif choice == "4":
            break
        else:
            print("无效选择，请重试")
            continue
        
        print()
        print("=" * 60)
    
    # 保存配置
    generator._save_config()
    print(f"配置文件已更新: {generator.config_path}")
    print("程序结束")

if __name__ == "__main__":
    main()
